{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp base_params\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from nbdev.showdoc import show_doc\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Params\n",
    "\n",
    "`BaseParams` is the major object to control the whole modeling process. It is supposed to be accessable anywhere. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# export\n",
    "\n",
    "import fcntl\n",
    "import json\n",
    "import os\n",
    "import re\n",
    "import time\n",
    "from collections import defaultdict\n",
    "from distutils.dir_util import copy_tree\n",
    "from typing import Any, Callable, Dict, List, Tuple, Union\n",
    "\n",
    "import tensorflow as tf\n",
    "from loguru import logger\n",
    "from m3tl.mtl_model.base import BasicMTL\n",
    "from m3tl.special_tokens import BOS_TOKEN, EOS_TOKEN, TRAIN\n",
    "from m3tl.utils import (create_path, get_phase, load_transformer_config,\n",
    "                        load_transformer_tokenizer)\n",
    "\n",
    "\n",
    "class BaseParams():\n",
    "    # pylint: disable=attribute-defined-outside-init\n",
    "    def __init__(self):\n",
    "\n",
    "        ################ transformers params #################\n",
    "        self.transformer_model_name = 'bert-base-chinese'\n",
    "        self.transformer_tokenizer_name = 'bert-base-chinese'\n",
    "        self.transformer_config_name = 'bert-base-chinese'\n",
    "        self.transformer_model_loading = 'TFBertModel'\n",
    "        self.transformer_config_loading = 'BertConfig'\n",
    "        self.transformer_tokenizer_loading = 'BertTokenizer'\n",
    "        self.transformer_decoder_model_name = None\n",
    "        self.transformer_decoder_config_name = None\n",
    "        self.transformer_decoder_tokenizer_name = None\n",
    "        # self.transformer_decoder_model_name = \"hfl/chinese-xlnet-base\"\n",
    "        # self.transformer_decoder_config_name = \"hfl/chinese-xlnet-base\"\n",
    "        # self.transformer_decoder_tokenizer_name = \"hfl/chinese-xlnet-base\"\n",
    "        self.transformer_decoder_model_loading = 'TFAutoModel'\n",
    "        self.transformer_decoder_config_loading = 'AutoConfig'\n",
    "        self.transformer_decoder_tokenizer_loading = 'AutoTokenizer'\n",
    "        ################ transformers params end ###############\n",
    "\n",
    "        ################ multi-model configuration #############\n",
    "        # multimodal params\n",
    "        self.modal_type_id = {\n",
    "            'text': 0,\n",
    "            'image': 1,\n",
    "            'others': 2\n",
    "        }\n",
    "        self.enable_modal_type = False\n",
    "        ############### multi-modal configuration end ##########\n",
    "        # bert config\n",
    "        self.init_checkpoint = ''\n",
    "\n",
    "        ############### registration information #############\n",
    "        # problem type\n",
    "        self.problem_type_list = []\n",
    "        self.predefined_problem_type = []\n",
    "        self.get_or_make_label_encoder_fn_dict: Dict[str, Callable] = {}\n",
    "        self.label_handling_fn: Dict[str, Callable] = {}\n",
    "        self.problem_type_desc : Dict[str, str] = {}\n",
    "        self.top_layer = {}\n",
    "        self.num_classes = {}\n",
    "        self.problem_info = defaultdict(dict)\n",
    "        self.problem_type = {}\n",
    "\n",
    "        # MTL models\n",
    "        self.available_mtl_model = {}\n",
    "        self.mtl_model = {}\n",
    "        self.register_mtl_model('basic', BasicMTL, include_top=False)\n",
    "        self.assign_mtl_model('basic')\n",
    "\n",
    "        # problem sampling strategy\n",
    "        self.available_problem_sampling_strategy = {}\n",
    "        self.problem_sampling_strategy = {}\n",
    "\n",
    "        # loss combination strategy\n",
    "        self.available_loss_combination_strategy = {}\n",
    "        self.loss_combination_strategy = {}\n",
    "\n",
    "        # gradient surgery\n",
    "        self.available_gradient_surgery = {}\n",
    "        self.gradient_surgery = {}\n",
    "\n",
    "        # embedding layer\n",
    "        self.available_embedding_layer = {}\n",
    "        self.embedding_layer = {}\n",
    "        ############### registration information end ###########\n",
    "\n",
    "        ############### logging control #############\n",
    "        self.detail_log = True\n",
    "        self.print_data = True\n",
    "        ############### logging control end #############\n",
    "\n",
    "        ############### model hparam #############\n",
    "        self.init_lr = 2e-5\n",
    "        self.dropout = 0.1\n",
    "        self.dropout_keep_prob = 1 - self.dropout\n",
    "        self.max_seq_len = 256\n",
    "        self.use_one_hot_embeddings = True\n",
    "        self.label_smoothing = 0.0\n",
    "        self.crf = False\n",
    "        self.bert_num_hidden_layer = 12\n",
    "        self.hidden_dense = False\n",
    "        # threshold to calculate metrics for multi_cls\n",
    "        self.multi_cls_threshold = 0.5\n",
    "        self.multi_cls_positive_weight = 1.0\n",
    "        self.custom_pooled_hidden_size = 0\n",
    "        self.share_embedding = True\n",
    "        self.output_problem_pred = True\n",
    "        self.output_body_seq_hidden = False\n",
    "        self.output_body_pooled_hidden = False\n",
    "        self.output_mtl_model_hidden = False\n",
    "        # DuplicateAugMultimodalEmbedding parameters\n",
    "        self.duplicate_data_aug_problems: Union[List[str], str] = None\n",
    "        # Contrastive Learning parameters\n",
    "        self.contrastive_learning_model_name: str = None\n",
    "        self.simcse_pooler = 'pooled'\n",
    "        ############### model hparam end #############\n",
    "\n",
    "        ############### training config #############\n",
    "        self.num_cpus = 1\n",
    "        self.preprocess_buffer = 100000\n",
    "        self.example_per_file = 100000\n",
    "        self.decode_vocab_file = None\n",
    "        self.batch_size = 32\n",
    "        self.train_epoch = 15\n",
    "        self.freeze_step = 0\n",
    "        self.prefetch = 5000\n",
    "        self.dynamic_padding = True\n",
    "        self.bucket_batch_sizes = [32, 32, 32, 16]\n",
    "        self.bucket_boundaries = [30, 64, 128]\n",
    "        self.shuffle_buffer = 200000\n",
    "        self.train_steps = 0\n",
    "        self.num_warmup_steps = 0\n",
    "        self.use_horovod = False\n",
    "        ############### training config end #############\n",
    "\n",
    "        ############### runtime assigned_details #############\n",
    "        self.run_problem_list = []\n",
    "        self.assigned_details = ()\n",
    "        self.problem_assigned = False\n",
    "        self.predicting = False\n",
    "        self.problem_list = []\n",
    "        self.problem_chunk = []\n",
    "        self.problem_str = None\n",
    "        ############### runtime assigned details end #############\n",
    "\n",
    "        ############### assign details #############\n",
    "        self.model_dir = None\n",
    "        self.ckpt_dir = ''\n",
    "        self.params_path = ''\n",
    "        self.bert_config = None\n",
    "        self.bert_config_dict = None\n",
    "        self.bert_decoder_config = None\n",
    "        self.bert_decoder_config_dict = None\n",
    "        self.init_weight_from_huggingface = True\n",
    "        self.vocab_size = None\n",
    "        self.decoder_vocab_size = None\n",
    "        self.bos_id = None\n",
    "        self.eos_id = None\n",
    "\n",
    "        ############### assign details end #############\n",
    "\n",
    "        self.tmp_file_dir = 'tmp'\n",
    "        # get generator function for each problem\n",
    "        self.read_data_fn = {}\n",
    "        self.problem_assigned = False\n",
    "\n",
    "        # pyspark related\n",
    "        self.pyspark_output_path = None\n",
    "\n",
    "        ############### legacy parameters #############\n",
    "        # seq2seq\n",
    "        self.decoder_num_hidden_layers = 3\n",
    "        self.beam_size = 10\n",
    "        self.init_decoder_from_encoder = False\n",
    "        self.beam_search_alpha = 0.6\n",
    "        self.decode_max_seq_len = 90\n",
    "\n",
    "        # pretrain hparm\n",
    "        self.dupe_factor = 10\n",
    "        self.short_seq_prob = 0.1\n",
    "        self.masked_lm_prob = 0.15\n",
    "        self.max_predictions_per_seq = 20\n",
    "        self.mask_lm_hidden_size = 768\n",
    "        self.mask_lm_hidden_act = 'gelu'\n",
    "        self.mask_lm_initializer_range = 0.02\n",
    "        ############### legacy parameters end #############\n",
    "\n",
    "    def get(self, key, default=None):\n",
    "        return getattr(self, key, default)\n",
    "\n",
    "    def to_json(self):\n",
    "        \"\"\"Save the params as json files. Please note that processing_fn is not saved.\n",
    "        \"\"\"\n",
    "        dump_dict = {}\n",
    "        for att_name, att in vars(self).items():\n",
    "            try:\n",
    "                json.dumps(att)\n",
    "                dump_dict[att_name] = att\n",
    "            except TypeError:\n",
    "                pass\n",
    "\n",
    "        with open(self.params_path, 'w', encoding='utf8') as f:\n",
    "            json.dump(dump_dict, f, indent=4)\n",
    "\n",
    "    def from_json(self, json_path: str = None):\n",
    "        \"\"\"Load json file as params.\n",
    "\n",
    "        json_path could not be None if the problem is not assigned to params\n",
    "\n",
    "        Args:\n",
    "            json_path (str, optional): Path to json file. Defaults to None.\n",
    "\n",
    "        Raises:\n",
    "            AttributeError\n",
    "        \"\"\"\n",
    "        try:\n",
    "            params_path = json_path if json_path is not None else self.params_path\n",
    "        except AttributeError:\n",
    "            AttributeError(\n",
    "                'Either json_path should not be None or problem is assigned.')\n",
    "\n",
    "        if self.problem_assigned:\n",
    "            assign_details = self.assigned_details\n",
    "        else:\n",
    "            assign_details = None\n",
    "\n",
    "        with open(params_path, 'r', encoding='utf8') as f:\n",
    "            dump_dict = json.load(f)\n",
    "        for att in dump_dict:\n",
    "            setattr(self, att, dump_dict[att])\n",
    "        self.bert_config = load_transformer_config(\n",
    "            self.bert_config_dict, self.transformer_config_loading)\n",
    "        if self.bert_decoder_config_dict:\n",
    "            self.bert_decoder_config = load_transformer_config(\n",
    "                self.bert_decoder_config_dict, self.transformer_decoder_config_loading\n",
    "            )\n",
    "        if assign_details:\n",
    "            self.assign_problem(*assign_details)\n",
    "\n",
    "    def parse_problem_string(self, flag_string: str) -> Tuple[List[str], List[List[str]]]:\n",
    "\n",
    "        self.problem_str = flag_string\n",
    "        problem_list = sorted(re.split(r'[&|]', flag_string))\n",
    "\n",
    "        for problem in problem_list:\n",
    "            if problem not in self.problem_type:\n",
    "                raise KeyError(\"problem {} is not registered.\".format(problem))\n",
    "\n",
    "        # Parse problem string\n",
    "        self.run_problem_list = []\n",
    "        problem_chunk = []\n",
    "        for flag_chunk in flag_string.split('|'):\n",
    "\n",
    "            if '&' not in flag_chunk:\n",
    "                problem_type = {}\n",
    "                problem_type[flag_chunk] = self.problem_type[flag_chunk]\n",
    "                self.run_problem_list.append(problem_type)\n",
    "                problem_chunk.append([flag_chunk])\n",
    "            else:\n",
    "                problem_type = {}\n",
    "                problem_chunk.append([])\n",
    "                for problem in flag_chunk.split('&'):\n",
    "                    problem_type[problem] = self.problem_type[problem]\n",
    "                    problem_chunk[-1].append(problem)\n",
    "                self.run_problem_list.append(problem_type)\n",
    "\n",
    "        return problem_list, problem_chunk\n",
    "\n",
    "    @staticmethod\n",
    "    def _copy_or_wait(from_path, to_path, do_copy=True, max_wait=5, replace=False):\n",
    "        if os.path.exists(to_path) and replace is False:\n",
    "            return\n",
    "        if do_copy:\n",
    "            copy_tree(from_path, to_path)\n",
    "            return\n",
    "        else:\n",
    "            for _ in range(max_wait):\n",
    "                if not os.path.exists(to_path):\n",
    "                    time.sleep(10)\n",
    "                else:\n",
    "                    return\n",
    "\n",
    "    def prepare_dir(self, base_dir: str, dir_name: str, problem_list: List[str], copy_dir=True):\n",
    "        \"\"\"prepare model checkpoint dir. this function will copy or save transformers' configs\n",
    "        and tokenizers to params.ckpt_dir\n",
    "\n",
    "        Args:\n",
    "            base_dir (str): base_dir of params.ckpt_dir. same as os.path.dirname(params.ckpt_dir). bad naming\n",
    "            dir_name (str): dir_name, same as os.path.basename(params.ckpt_dir). bad naming\n",
    "            problem_list (List[str]): [description]\n",
    "        \"\"\"\n",
    "        base = base_dir if base_dir is not None else 'models'\n",
    "\n",
    "        dir_name = dir_name if dir_name is not None else '_'.join(\n",
    "            problem_list)+'_ckpt'\n",
    "        self.ckpt_dir = os.path.join(base, dir_name)\n",
    "\n",
    "        # we need to make sure all configs, tokenizers are in ckpt_dir\n",
    "        # configs\n",
    "        from_config_path = os.path.join(self.init_checkpoint,\n",
    "                                        'bert_config')\n",
    "        from_decoder_config_path = os.path.join(self.init_checkpoint,\n",
    "                                                'bert_decoder_config')\n",
    "        to_config_path = os.path.join(self.ckpt_dir, 'bert_config')\n",
    "        to_decoder_config_path = os.path.join(\n",
    "            self.ckpt_dir, 'bert_decoder_config')\n",
    "\n",
    "        # tokenizers\n",
    "        from_tokenizer_path = os.path.join(self.init_checkpoint, 'tokenizer')\n",
    "        to_tokenizer_path = os.path.join(self.ckpt_dir, 'tokenizer')\n",
    "\n",
    "        from_decoder_tokenizer_path = os.path.join(\n",
    "            self.init_checkpoint, 'decoder_tokenizer')\n",
    "        to_decoder_tokenizer_path = os.path.join(\n",
    "            self.ckpt_dir, 'decoder_tokenizer')\n",
    "\n",
    "        self.params_path = os.path.join(self.ckpt_dir, 'params.json')\n",
    "\n",
    "        if not self.predicting:\n",
    "            create_path(self.ckpt_dir)\n",
    "\n",
    "            # two ways to init model\n",
    "            # 1. init from TF checkpoint dir created by m3tl.\n",
    "            # 2. init from huggingface checkpoint.\n",
    "\n",
    "            # bert config exists, init from existing config\n",
    "            if os.path.exists(from_config_path):\n",
    "                # copy config\n",
    "                self._copy_or_wait(\n",
    "                    from_config_path, to_config_path, do_copy=copy_dir)\n",
    "                self.bert_config = load_transformer_config(\n",
    "                    to_config_path, self.transformer_config_loading)\n",
    "\n",
    "                # copy tokenizer\n",
    "                self._copy_or_wait(from_tokenizer_path,\n",
    "                                   to_tokenizer_path, do_copy=copy_dir)\n",
    "\n",
    "                # copy decoder config\n",
    "                if os.path.exists(from_decoder_config_path):\n",
    "                    self._copy_or_wait(\n",
    "                        from_decoder_config_path, to_decoder_config_path, do_copy=copy_dir)\n",
    "                    self.bert_decoder_config = load_transformer_config(\n",
    "                        from_decoder_config_path, self.transformer_decoder_config_loading\n",
    "                    )\n",
    "                    self.bert_decoder_config_dict = self.bert_decoder_config.to_dict()\n",
    "                # copy decoder tokenizer\n",
    "                if os.path.exists(from_decoder_tokenizer_path):\n",
    "                    self._copy_or_wait(\n",
    "                        from_decoder_tokenizer_path, to_decoder_tokenizer_path, do_copy=copy_dir)\n",
    "\n",
    "                self.init_weight_from_huggingface = False\n",
    "            else:\n",
    "                # load config from huggingface\n",
    "                logger.warning(\n",
    "                    '{} not exists. will load model from huggingface checkpoint.', from_config_path)\n",
    "                # get or download config\n",
    "                self.init_weight_from_huggingface = True\n",
    "                self.bert_config = load_transformer_config(\n",
    "                    self.transformer_config_name, self.transformer_config_loading)\n",
    "                self.bert_config.save_pretrained(to_config_path)\n",
    "\n",
    "                # save tokenizer\n",
    "                tokenizer = load_transformer_tokenizer(\n",
    "                    self.transformer_tokenizer_name, self.transformer_tokenizer_loading)\n",
    "                tokenizer.save_pretrained(to_tokenizer_path)\n",
    "                # save_pretrained method of tokenizer saves the config as tokenizer_config.json, which will cause\n",
    "                # OSError if use tokenizer.from_pretrained directly. we need to manually rename the json file\n",
    "                try:\n",
    "                    os.rename(os.path.join(to_tokenizer_path, 'tokenizer_config.json'), os.path.join(\n",
    "                        to_tokenizer_path, 'config.json'))\n",
    "                except:\n",
    "                    pass\n",
    "\n",
    "                # if decoder is specified\n",
    "                if self.transformer_decoder_model_name:\n",
    "                    self.bert_decoder_config = load_transformer_config(\n",
    "                        self.transformer_decoder_config_name, self.transformer_decoder_config_loading\n",
    "                    )\n",
    "                    self.bert_decoder_config_dict = self.bert_decoder_config.to_dict()\n",
    "                    self.bert_decoder_config.save_pretrained(\n",
    "                        to_decoder_config_path)\n",
    "                    decoder_tokenizer = load_transformer_tokenizer(\n",
    "                        self.transformer_decoder_tokenizer_name, self.transformer_decoder_tokenizer_loading)\n",
    "                    decoder_tokenizer.save_pretrained(\n",
    "                        to_decoder_tokenizer_path)\n",
    "                    try:\n",
    "                        os.rename(os.path.join(to_decoder_tokenizer_path, 'tokenizer_config.json'), os.path.join(\n",
    "                            to_decoder_tokenizer_path, 'config.json'))\n",
    "                    except:\n",
    "                        pass\n",
    "        else:\n",
    "            self.bert_config = load_transformer_config(to_config_path)\n",
    "            if os.path.exists(to_decoder_config_path):\n",
    "                self.bert_decoder_config = load_transformer_config(\n",
    "                    to_decoder_config_path)\n",
    "            self.init_weight_from_huggingface = False\n",
    "\n",
    "        self.transformer_config_name = to_config_path\n",
    "        # set value if and only if decoder is assigned\n",
    "        self.transformer_decoder_config_name = to_decoder_config_path if self.transformer_decoder_config_name is not None else None\n",
    "        self.transformer_tokenizer_name = to_tokenizer_path\n",
    "        # set value if and only if decoder is assigned\n",
    "        self.transformer_decoder_tokenizer_name = to_decoder_tokenizer_path if self.transformer_decoder_tokenizer_name is not None else None\n",
    "\n",
    "        self.bert_config_dict = self.bert_config.to_dict()\n",
    "\n",
    "        tokenizer = load_transformer_tokenizer(\n",
    "            self.transformer_tokenizer_name, self.transformer_tokenizer_loading)\n",
    "        self.vocab_size = tokenizer.vocab_size\n",
    "        if self.transformer_decoder_tokenizer_name:\n",
    "            decoder_tokenizer = load_transformer_tokenizer(\n",
    "                self.transformer_decoder_tokenizer_name,\n",
    "                self.transformer_decoder_tokenizer_loading\n",
    "            )\n",
    "\n",
    "            # if set bos and eos\n",
    "            if decoder_tokenizer.bos_token is None:\n",
    "                decoder_tokenizer.add_special_tokens({'bos_token': BOS_TOKEN})\n",
    "\n",
    "            if decoder_tokenizer.eos_token is None:\n",
    "                decoder_tokenizer.add_special_tokens({'eos_token': EOS_TOKEN})\n",
    "\n",
    "            # overwrite tokenizer\n",
    "            decoder_tokenizer.save_pretrained(to_decoder_tokenizer_path)\n",
    "\n",
    "            self.decoder_vocab_size = decoder_tokenizer.vocab_size\n",
    "            self.bos_id = decoder_tokenizer.bos_token_id\n",
    "            self.eos_id = decoder_tokenizer.eos_token_id\n",
    "\n",
    "    def get_problem_type(self, problem: str) -> str:\n",
    "        return self.problem_type[problem]\n",
    "\n",
    "    def update_train_steps(self, train_steps_per_epoch: int, epoch: int = None, warmup_ratio=0.1) -> None:\n",
    "        \"\"\"If the batch_size is dynamic, we have to loop through the tf.data.Dataset\n",
    "        to get the accurate number of training steps. In this case, we need a function to\n",
    "        update the train_steps which will be used to calculate learning rate schedule.\n",
    "\n",
    "        WARNING: updating should be called before the model is compiled!\n",
    "\n",
    "        Args:\n",
    "            train_steps (int): new number of train_steps\n",
    "        \"\"\"\n",
    "        if epoch:\n",
    "            train_steps = train_steps_per_epoch * epoch\n",
    "        else:\n",
    "            train_steps = train_steps_per_epoch * self.train_epoch\n",
    "\n",
    "        logger.critical('Updating train_steps to {}'.format(train_steps))\n",
    "\n",
    "        self.train_steps = train_steps\n",
    "        self.train_steps_per_epoch = train_steps_per_epoch\n",
    "        self.num_warmup_steps = int(self.train_steps * warmup_ratio)\n",
    "\n",
    "    def get_problem_chunk(self, problem: str = None, as_str=True) -> Union[List[str], List[List[str]]]:\n",
    "\n",
    "        # get chunk of the problem\n",
    "        if problem is not None:\n",
    "            all_problem_chunk_as_str = self.get_problem_chunk(as_str=True)\n",
    "            # if problem itself is actually a chunk\n",
    "            if problem in all_problem_chunk_as_str:\n",
    "                return [problem]\n",
    "            problem_chunk_to_search = [\n",
    "                c for c in self.problem_chunk if problem in c]\n",
    "        else:\n",
    "            problem_chunk_to_search = self.problem_chunk\n",
    "\n",
    "        if len(problem_chunk_to_search) == 0:\n",
    "            raise ValueError('problem {} not in problem chunk {}'.format(\n",
    "                problem, self.problem_chunk))\n",
    "\n",
    "        if as_str:\n",
    "            res_list = []\n",
    "            for problem_list in problem_chunk_to_search:\n",
    "                res_list.append('_'.join(sorted(problem_list)))\n",
    "            return res_list\n",
    "        else:\n",
    "            return problem_chunk_to_search\n",
    "\n",
    "    def register_problem_type(self,\n",
    "                              problem_type: str,\n",
    "                              top_layer: tf.keras.Model = None,\n",
    "                              label_handling_fn: Callable = None,\n",
    "                              get_or_make_label_encoder_fn: Callable = None,\n",
    "                              inherit_from: str = None,\n",
    "                              description: str = None):\n",
    "        self.problem_type_list.append(problem_type)\n",
    "\n",
    "        if inherit_from is not None:\n",
    "            assert inherit_from in self.problem_type_list, 'try to inherit from {0}, but {0} is not registered.'.format(\n",
    "                inherit_from)\n",
    "\n",
    "        if get_or_make_label_encoder_fn is None and inherit_from is not None:\n",
    "            self.get_or_make_label_encoder_fn_dict[problem_type] = self.get_or_make_label_encoder_fn_dict[inherit_from]\n",
    "        else:\n",
    "            if not isinstance(get_or_make_label_encoder_fn, Callable):\n",
    "                raise TypeError(\n",
    "                    'get_or_make_label_encoder_fn of problem type {} should be callable'.format(problem_type))\n",
    "            self.get_or_make_label_encoder_fn_dict[problem_type] = get_or_make_label_encoder_fn\n",
    "\n",
    "        if top_layer is None and inherit_from is not None:\n",
    "            self.top_layer[problem_type] = self.top_layer[inherit_from]\n",
    "        else:\n",
    "            # if not isinstance(top_layer, tf.keras.Model) and not isinstance(top_layer, tf.keras.layers.Layer):\n",
    "            #     raise TypeError(\n",
    "            #         'top_layer of problem {} should be `tf.keras.Model` or `tf.keras.layers.Layer`'.format(problem_type))\n",
    "            self.top_layer[problem_type] = top_layer\n",
    "\n",
    "        if label_handling_fn is None and inherit_from is not None:\n",
    "            self.label_handling_fn[problem_type] = self.label_handling_fn[inherit_from]\n",
    "        else:\n",
    "            if not isinstance(label_handling_fn, Callable):\n",
    "                raise TypeError(\n",
    "                    'label_handling_fn of problem type {} should be callable'.format(problem_type))\n",
    "            self.label_handling_fn[problem_type] = label_handling_fn\n",
    "        \n",
    "        self.problem_type_desc[problem_type] = description\n",
    "\n",
    "    def list_available_problem_types(self) -> List[str]:\n",
    "        return self.problem_type_list\n",
    "\n",
    "    def register_problem(self, problem_name: str, problem_type='cls', processing_fn: Callable = None):\n",
    "\n",
    "        if problem_type not in self.problem_type_list:\n",
    "            raise ValueError('Provided problem type not valid, expect {0}, got {1}'.format(\n",
    "                self.problem_type_list,\n",
    "                problem_type))\n",
    "\n",
    "        self.problem_type[problem_name] = problem_type\n",
    "        self.read_data_fn[problem_name] = processing_fn\n",
    "\n",
    "    def register_multiple_problems(self, problem_type_dict: Dict[str, str], processing_fn_dict: Dict[str, Callable] = None):\n",
    "        # add new problem to params if problem_type_dict and processing_fn_dict provided\n",
    "        for new_problem, problem_type in problem_type_dict.items():\n",
    "            logger.info('Adding new problem {0}, problem type: {1}'.format(\n",
    "                new_problem, problem_type_dict[new_problem]))\n",
    "            if processing_fn_dict:\n",
    "                new_problem_processing_fn = processing_fn_dict[new_problem]\n",
    "            else:\n",
    "                new_problem_processing_fn = None\n",
    "            self.register_problem(\n",
    "                problem_name=new_problem, problem_type=problem_type, processing_fn=new_problem_processing_fn)\n",
    "\n",
    "    @staticmethod\n",
    "    def _append_info_to_file(path: str, problem: str, info_name: str, info: Any):\n",
    "        if ':' in info_name:\n",
    "            raise ValueError('info_name cannot contain :')\n",
    "        if not os.path.exists(path):\n",
    "            os.makedirs(os.path.dirname(path), exist_ok=True)\n",
    "        with open(path, 'a', encoding='utf8') as f:\n",
    "            fcntl.flock(f, fcntl.LOCK_EX)\n",
    "            info_str = json.dumps(info)\n",
    "            f.write('{}:{}:{}\\n'.format(problem, info_name, info_str))\n",
    "            fcntl.flock(f, fcntl.LOCK_UN)\n",
    "\n",
    "    @staticmethod\n",
    "    def _read_info_file(path: str) -> dict:\n",
    "        res = defaultdict(dict)\n",
    "        if not os.path.exists(path):\n",
    "            return res\n",
    "        with open(path, 'r', encoding='utf8') as f:\n",
    "            for line in f:\n",
    "                problem, key, val = line.split(':', 2)\n",
    "                val = json.loads(val)\n",
    "                res[problem][key] = val\n",
    "        return res\n",
    "\n",
    "    def merge_problem_info_file(self, path: str):\n",
    "        \"\"\"merge a problem info file\"\"\"\n",
    "        income_problem_info = BaseParams._read_info_file(path)\n",
    "        for problem in income_problem_info:\n",
    "            for info_name, info in income_problem_info[problem].items():\n",
    "                self.set_problem_info(\n",
    "                    problem=problem, info_name=info_name, info=info)\n",
    "\n",
    "    def set_problem_info(self, problem: str, info_name: str, info: Any):\n",
    "        \"\"\"api to save problem-specific info and can be access anywhere\"\"\"\n",
    "        if not problem or get_phase() != TRAIN:\n",
    "            return\n",
    "\n",
    "        # if json exists, load json\n",
    "        problem_info_json_path = self.get_problem_info_path(problem=problem)\n",
    "\n",
    "        try:\n",
    "            old_info = self.get_problem_info(\n",
    "                problem=problem, info_name=info_name)\n",
    "            if old_info != info:\n",
    "                raise KeyError\n",
    "        except KeyError:\n",
    "            self._append_info_to_file(\n",
    "                path=problem_info_json_path, problem=problem, info_name=info_name, info=info)\n",
    "            self.problem_info.update(\n",
    "                self._read_info_file(problem_info_json_path))\n",
    "\n",
    "    def get_problem_info_path(self, problem) -> str:\n",
    "        problem_chunk = self.get_problem_chunk(problem=problem, as_str=True)[0]\n",
    "        return os.path.join(self.tmp_file_dir, problem_chunk, 'problem_info.txt')\n",
    "\n",
    "    def get_problem_info(self, problem: str = None, info_name: str = None) -> Any:\n",
    "        if problem is None:\n",
    "            # load all jsons\n",
    "            for problem_chunk in self.get_problem_chunk():\n",
    "                self.get_problem_info(problem=problem_chunk)\n",
    "            for problem in self.problem_list:\n",
    "                self.get_problem_info(problem=problem)\n",
    "            return self.problem_info\n",
    "\n",
    "        problem_info_json_path = self.get_problem_info_path(problem=problem)\n",
    "        if os.path.exists(problem_info_json_path):\n",
    "            self.problem_info.update(\n",
    "                self._read_info_file(problem_info_json_path))\n",
    "\n",
    "        if problem not in self.problem_info:\n",
    "            err_str = (\"cannot get {0}'s problem info, seems you have not created tfrecord for problem. \\n\"\n",
    "                       \"please make sure either: \\n\"\n",
    "                       \"    1. path {1} exists. this is usually created when creating tfrecord for problem. \\n\"\n",
    "                       \"    2. {0} exists in 'problem_info' key of {2}. \".format(problem, problem_info_json_path, self.params_path))\n",
    "            raise KeyError(err_str)\n",
    "\n",
    "        if info_name is None:\n",
    "            return self.problem_info[problem]\n",
    "\n",
    "        if info_name not in self.problem_info[problem]:\n",
    "            raise KeyError(\n",
    "                \"problem {} dose not have info {}\".format(problem, info_name))\n",
    "\n",
    "        return self.problem_info[problem][info_name]\n",
    "\n",
    "    def assign_problem(self,\n",
    "                       flag_string: str,\n",
    "                       model_dir: str = None,\n",
    "                       base_dir: str = None,\n",
    "                       dir_name: str = None,\n",
    "                       predicting=False):\n",
    "\n",
    "        if model_dir:\n",
    "            base_dir, dir_name = os.path.split(model_dir)\n",
    "        elif base_dir or dir_name:\n",
    "            logger.warning(\n",
    "                \"base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\")\n",
    "\n",
    "        self.assigned_details = (\n",
    "            flag_string, model_dir, base_dir, dir_name, predicting)\n",
    "        self.problem_assigned = True\n",
    "        self.predicting = predicting\n",
    "\n",
    "        self.problem_list, self.problem_chunk = self.parse_problem_string(\n",
    "            flag_string)\n",
    "\n",
    "        # create dir and get vocab, config\n",
    "        if self.use_horovod:\n",
    "            import horovod.tensorflow.keras as hvd\n",
    "            if hvd.rank() == 0:\n",
    "                self.prepare_dir(base_dir, dir_name, self.problem_list)\n",
    "            else:\n",
    "                self.prepare_dir(base_dir, dir_name,\n",
    "                                 self.problem_list, copy_dir=False)\n",
    "\n",
    "            self.lr = self.init_lr * hvd.size()\n",
    "        else:\n",
    "            self.prepare_dir(base_dir, dir_name, self.problem_list)\n",
    "            self.lr = self.init_lr * \\\n",
    "                len(tf.config.experimental.list_physical_devices('GPU'))\n",
    "\n",
    "    def list_available_problems(self) -> Dict[str, str]:\n",
    "        return self.problem_type\n",
    "\n",
    "    def assign_data_sampling_strategy(self,\n",
    "                                      sampling_strategy_name='data_balanced',\n",
    "                                      sampling_strategy_fn: Callable = None):\n",
    "        self.problem_sampling_strategy['name'] = sampling_strategy_name\n",
    "        self.problem_sampling_strategy['fn'] = sampling_strategy_fn\n",
    "\n",
    "    def calculate_data_sampling_prob(self) -> Dict[str, float]:\n",
    "\n",
    "        problem_chunk_data_num = defaultdict(float)\n",
    "        sampling_strategy_name = self.problem_sampling_strategy['name']\n",
    "        if sampling_strategy_name == 'data_balanced':\n",
    "            problem_chunk = self.get_problem_chunk(as_str=True)\n",
    "            for problem_list in problem_chunk:\n",
    "                problem_chunk_data_num[problem_list] += self.get_problem_info(\n",
    "                    problem=problem_list, info_name='data_num')\n",
    "        elif sampling_strategy_name == 'problem_balanced':\n",
    "            problem_chunk = self.get_problem_chunk(as_str=True)\n",
    "            for str_per_chunk in problem_chunk:\n",
    "                problem_chunk_data_num[str_per_chunk] = 1\n",
    "        else:\n",
    "            raise ValueError(\n",
    "                'sampling strategy {} is not implemented by default. '\n",
    "                'please provide sampling_strategy_fn.'.format(sampling_strategy_name))\n",
    "\n",
    "        # devided by sum to get sampling prob\n",
    "        sum_across_problems = sum(\n",
    "            [v for _, v in problem_chunk_data_num.items()])\n",
    "        self.problem_sampling_weight_dict = {\n",
    "            k: v / sum_across_problems for k, v in problem_chunk_data_num.items()}\n",
    "        return self.problem_sampling_weight_dict\n",
    "\n",
    "    def list_available_mtl_setup(self) -> Dict[str, list]:\n",
    "        ret_dict = {}\n",
    "        ret_dict['available_mtl_model'] = list(self.available_mtl_model.keys())\n",
    "        ret_dict['available_problem_sampling_strategy'] = list(\n",
    "            self.available_problem_sampling_strategy.keys())\n",
    "        ret_dict['available_loss_combination_strategy'] = list(\n",
    "            self.available_loss_combination_strategy.keys())\n",
    "        ret_dict['available_gradient_surgery'] = list(\n",
    "            self.available_gradient_surgery)\n",
    "        return ret_dict\n",
    "\n",
    "    def register_mtl_model(\n",
    "            self,\n",
    "            model_name: str,\n",
    "            model: tf.keras.Model,\n",
    "            include_top: bool,\n",
    "            extra_info: dict = None):\n",
    "        self.available_mtl_model[model_name] = {}\n",
    "        self.available_mtl_model[model_name]['model'] = model\n",
    "        self.available_mtl_model[model_name]['include_top'] = include_top\n",
    "        self.available_mtl_model[model_name]['extra_info'] = extra_info\n",
    "        self.available_mtl_model[model_name]['model_name'] = model_name\n",
    "\n",
    "    def assign_mtl_model(self, model_name: str):\n",
    "        if model_name not in self.available_mtl_model:\n",
    "            raise KeyError('MTL model {} is not registered. Available MTL model: {}'.format(\n",
    "                model_name, self.available_mtl_model.keys()))\n",
    "        self.mtl_model = self.available_mtl_model[model_name]\n",
    "\n",
    "    def register_loss_combination_strategy(\n",
    "            self,\n",
    "            name: str,\n",
    "            model: tf.keras.Model,\n",
    "            callbacks: List[tf.keras.callbacks.Callback] = None,\n",
    "            extra_info: dict = None):\n",
    "        self.available_loss_combination_strategy[name] = {}\n",
    "        self.available_loss_combination_strategy[name]['model'] = model\n",
    "        self.available_loss_combination_strategy[name]['name'] = name\n",
    "        self.available_loss_combination_strategy[name]['callbacks'] = callbacks\n",
    "        self.available_loss_combination_strategy[name]['extra_info'] = extra_info\n",
    "\n",
    "    def assign_loss_combination_strategy(\n",
    "            self,\n",
    "            name: str):\n",
    "        if name not in self.available_loss_combination_strategy:\n",
    "            raise KeyError('Loss combination strategy {} is not registered. Available strategy: {}'.format(\n",
    "                name, self.available_loss_combination_strategy.keys()))\n",
    "        self.loss_combination_strategy = self.available_loss_combination_strategy[name]\n",
    "\n",
    "    def register_embedding_layer(self, name: str, embedding_layer: tf.keras.Model, **kwargs):\n",
    "        self.available_embedding_layer[name] = {}\n",
    "        self.available_embedding_layer[name]['model'] = embedding_layer\n",
    "        self.available_embedding_layer[name]['name'] = name\n",
    "\n",
    "        for key, arg in kwargs.items():\n",
    "            self.available_embedding_layer[name][key] = arg\n",
    "\n",
    "    def assign_embedding_layer(self, name: str):\n",
    "        if name not in self.available_embedding_layer:\n",
    "            raise KeyError('Embedding layer {} is not registered. Available embedding_layer: {}'.format(\n",
    "                name, self.available_embedding_layer.keys()))\n",
    "        self.embedding_layer = self.available_embedding_layer[name]\n",
    "\n",
    "    def get_current_setup(self) -> Dict[str, str]:\n",
    "        setup_dict = {}\n",
    "        if self.problem_assigned:\n",
    "            setup_dict['assigned_problems'] = self.assigned_details[0]\n",
    "\n",
    "        setup_dict['mtl_model'] = self.mtl_model['model_name']\n",
    "        setup_dict['problem_sampling_strategy'] = self.problem_sampling_strategy.get(\n",
    "            'name', 'data_balanced')\n",
    "        setup_dict['loss_combination_strategy'] = self.loss_combination_strategy.get(\n",
    "            'name', 'sum')\n",
    "        setup_dict['gradient_surgery'] = self.gradient_surgery.get(\n",
    "            'name', 'average')\n",
    "\n",
    "        return setup_dict\n",
    "\n",
    "    def gather_mtl_callbacks(self):\n",
    "        callbacks = []\n",
    "        loss_callback = self.loss_combination_strategy['callbacks']\n",
    "        if loss_callback:\n",
    "            callbacks.append(loss_callback)\n",
    "        return callbacks\n",
    "\n",
    "\n",
    "class CRFParams(BaseParams):\n",
    "    def __init__(self):\n",
    "        super(CRFParams, self).__init__()\n",
    "        self.crf = True\n",
    "\n",
    "\n",
    "class StaticBatchParams(BaseParams):\n",
    "    def __init__(self):\n",
    "        super(StaticBatchParams, self).__init__()\n",
    "        self.dynamic_padding = False\n",
    "\n",
    "\n",
    "class DynamicBatchSizeParams(BaseParams):\n",
    "    def __init__(self):\n",
    "        super(DynamicBatchSizeParams, self).__init__()\n",
    "        self.bucket_batch_sizes = [128, 64, 32, 16]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-15 17:01:36.062 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-15 17:01:36.062 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-15 17:01:36.063 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-15 17:01:36.063 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-15 17:01:36.064 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_regression, problem type: regression\n",
      "2021-06-15 17:01:36.064 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "2021-06-15 17:01:36.065 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-15 17:01:36.066 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem fake_contrastive_learning, problem type: contrastive_learning\n",
      "2021-06-15 17:01:36.066 | WARNING  | m3tl.base_params:assign_problem:634 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-15 17:01:36.068 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-15 17:01:41.664 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "tb = TestBase()\n",
    "params = tb.params\n",
    "tmp_model_dir = tb.tmpckptdir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "# test inherit problem type\n",
    "params.register_problem_type(problem_type='clone_cls', inherit_from='cls')\n",
    "assert params.top_layer['cls'] == params.top_layer['clone_cls']\n",
    "\n",
    "# test correct error message of get_problem_info\n",
    "try:\n",
    "    params.get_problem_info(problem='abcd')\n",
    "except ValueError:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
